#!/Pub/Apps/bin/Rscript
# pkgs <- c("purrr", "argparse")
# .pkg_check <- \(pkgs) {
#     pkg_v <- \(pkg_name) {
#         packageVersion(pkg_name)
#     }
#     for (i in pkgs) {
#         tryCatch(pkg_v(i), error = \(e) {
#             message(paste0("Need pkg \033[4ma", i, "\033[24m, try to install it via utils::install.packages()\n"))
#             utils::install.packages(i, repos = "https://mirrors.bfsu.edu.cn/CRAN/")
#         })
#     }
# }
# .pkg_check(pkgs)
suppressPackageStartupMessages(library(purrr))
suppressPackageStartupMessages(library(argparse))
suppressPackageStartupMessages(library(tidyverse))

wrp <- "
    分子对接。
    测试参考：
    docking -c /Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/R/Convertional_Transcriptome/Molecular_Docking_cmd/config \
    -p /Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/R/Convertional_Transcriptome/Molecular_Docking_cmd/ALB_1ao6_nw.pdbqt \
    -od /Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/R/Convertional_Transcriptome/Molecular_Docking_cmd/test
    "
# create parser object-------
parser <- ArgumentParser(
    description = wrp,
    formatter_class = "argparse.RawTextHelpFormatter"
)

parser$add_argument("-c", "--config",
    action = "store",
    default = 'NULL', metavar = "",
    help = "config文件，参考文件路径 /Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/developing/分子对接并行计算/config"
)
parser$add_argument("-p", "--pdbqt",
    action = "store",
    default = 'NULL', metavar = "",
    help = "pdbqt文件，蛋白pdbqt文件路径，需要去水加氢观察是否有口袋"
)
parser$add_argument("-od", "--output_dir",
    action = "store",
    default = '/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/R/Convertional_Transcriptome/Molecular_Docking_cmd/test', metavar = "",
    help = "结果输出路径"
)
parser$add_argument("-d", "--drug",
    action = "store",
    default = 'NULL', metavar = "",
    help = "药物名称文件，可以筛选部分药物来做分子对接，全部药物名称见路径 /Pub/Users/wangyk/GEO_Data_ZhangGX/drug_names_for_md/drug_name.tsv"
)
parser$add_argument("-f", "--filter_res",
    action = "store",
    default = TRUE, metavar = "",
    help = "是否对结果进行筛选汇总"
)
parser$add_argument("-co","--core",
    action = "store",
    default = 4, metavar = "",
    help = "并行计算核心数，该参数决定多少个分子在同时进行对接，[default 4]"
)
parser$add_argument("-cp", "--cpu",
    action = "store",
    default = 4, metavar = "",
    help = "CPU数目，在单个分子对接过程中的cpu数目，是vina命令中的参数，默认4，该参数决定了单个分子对接的运行速度"
)

# parser$print_help()
args <- parser$parse_args()

jd <- file_test("-f", args$config)
jf <- file_test("-f", args$pdbqt)

if (jd + jf < 2) {
    cli::cli_alert_danger("config 或者pdbqt文件有问题 可能有问题，检查一下.")
    q(status = 1)
}

if (jd + jf == 2) {
    suppressMessages(library(tidyverse))
    suppressMessages(library(magrittr))
    library(furrr)

    od <- args$output_dir
    dir.create(od)
    config_file <- args$config
    setwd(dirname(config_file))

    pdbqt_files <- list.files("/Pub/Users/wangyk/project/Poroject/Molecular_ducking/pdbqt_file", pattern = "pdbqt", full.names = T)

    if(args$drug != 'NULL'){
        drug_filtered <- read_lines(args$drug)
        pdbqt_files <- map(drug_filtered,~ str_subset(pdbqt_files,.x)) %>% unlist()
    }

    drug_name <- str_split(pdbqt_files, "/") %>%
        map_chr(~ .x[9]) %>%
        str_remove("\\.pdbqt")

    ducking_res_od <- file.path(od, drug_name)
    plan("multicore", workers = as.integer(args$core))


    ducking <- \(x, y, z) {
        dir.create(y, showWarnings = F, recursive = T)
        system2(
            command = "/Pub/Users/wangyk/miniconda3/envs/Molecular_docking/bin/vina",
            # command = "/Pub/Users/wangyk/software/molecular_dockinig/vina/.pixi/envs/default/bin/vina",
            args = c(
                str_glue("--config {config_file}"),
                str_glue("--ligand {x}"),
                str_glue("--out {y}/out.pdbqt"),
                str_glue("--exhaustiveness 20"),
                str_glue("--cpu {args$cpu}"),
                str_glue("--energy_range 4"),
                str_glue("--num_modes 5"),
                str_glue("--receptor {args$pdbqt}"),
                str_glue("--log {y}/{z}_log.txt")
            )
        )
    }

    furrr::future_pwalk(list(pdbqt_files, ducking_res_od, drug_name), ducking)

    message(str_glue("————对接完成————"))
} else {
   cli::cli_alert_info("缺少 config或者pdbqt文件，如果od是已完成分析的结果路径，忽略该提示")
}


# print(args$filter_res)
# print(args$output_dir)
# args$output_dir <- "/Pub/Users/wangyk/project/p/zhangjuan_分子对接_20240412/res/ALB"
if (args$filter_res) {
    list.files(
        args$output_dir, "\\.txt",
        recursive = T, full.names = T
    ) -> file_dirs

    res <- map(seq_along(file_dirs), function(i) {
        if (str_detect(file_dirs[i], "config")) {
            return(0)
        }

        reading <- file_dirs[i] %>%
            read_lines(skip = 22, progress = F)

        if(length(reading) == 0){
            message(str_glue("-该文件对接结果为0，{file_dirs[i]} \n\n"))
            return(0)
        }

        reading[4:{
            which(reading == "Writing output ... done.") - 1
        }] %>%
            str_split(" ") -> a

        map(a, ~ which(.x != "")) -> index

        res <- map2(a, index, ~ .x[.y]) %>%
            map(., ~ .x[2]) %>%
            unlist() %>%
            as.numeric() %>%
            .[1]

        return(res)
    })


    names(res) <- seq_along(res)
    i <- sort(unlist(res)) %>% names() %>% as.numeric() %>% .[1:5]
    cli::cli_h3("\n结合能最低的top5对接log路径：")
    message( str_c(file_dirs[i],'\n'))
}

message(str_glue("————分析完成————"))



# log文件解读参考(AI问答)：
# 在AutoDock Vina的分子对接log文件中，以下列的含义如下：
# 1. mode:
#    - 代表对接模式的编号。Vina会生成多个对接模式（poses），每个模式对应一个可能的配体-受体结合构象。

# 2. affinity:
#    - 代表每个对接模式的结合自由能（binding affinity），单位是kcal/mol。数值越负，表示结合越稳定。

# 3. dist from rmsd l.b.:
#    - 代表当前对接模式与最优模式（affinity最负的模式）之间的RMSD（Root Mean Square Deviation）下限。RMSD是衡量两个构象之间差异的指标，单位是埃（Å）。下限表示最小可能的RMSD值。

# 4. best mode rmsd l.b.:
#    - 代表当前对接模式与所有模式中最优模式（affinity最负的模式）之间的RMSD下限。这个值用于比较当前模式与最优模式的相似性。

# ### 示例
# 假设log文件中有以下内容：
# ```
# mode | affinity | dist from best mode rmsd l.b. | dist from best mode rmsd u.b.
# 1    | -7.5     | 0.0                           | 0.0
# 2    | -7.3     | 2.1                           | 2.5
# 3    | -7.0     | 1.8                           | 2.2
# ```

# - mode 1 是最优模式，affinity为-7.5 kcal/mol，RMSD为0.0（与自身比较）。
# - mode 2 的affinity为-7.3 kcal/mol，与最优模式的RMSD下限为2.1 Å。
# - mode 3 的affinity为-7.0 kcal/mol，与最优模式的RMSD下限为1.8 Å。

# ### 总结
# 这些列帮助研究者评估不同对接模式的结合稳定性和相似性，从而选择最有可能的配体-受体结合构象。